import torch
import cv2
from config import MODEL_PATH, TARGET_CLASSES, CONF_THRESHOLD
from distance_utils import estimate_distance
from saver import ImageSaver
from logger import get_logger
import numpy as np

logger = get_logger('Detector')

class DroneDetector:
    def __init__(self):
        logger.info("Loading YOLO model...")
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=MODEL_PATH, force_reload=True)
        self.saver = ImageSaver()

    def detect_and_annotate(self, frame):
        results = self.model(frame)
        detections = []

        for *xyxy, conf, cls in results.xyxy[0]:
            label = self.model.names[int(cls)]
            if conf < CONF_THRESHOLD or label.lower() not in TARGET_CLASSES:
                continue

            x1, y1, x2, y2 = map(int, xyxy)
            pixel_width = x2 - x1
            distance = estimate_distance(pixel_width)

            detections.append({
                'label': label,
                'x1': x1, 'y1': y1, 'x2': x2, 'y2': y2,
                'confidence': float(conf),
                'distance_m': round(distance, 2)
            })

            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            text = f"{label} {distance:.2f}m"
            cv2.putText(frame, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

        if detections:
            self.saver.save(frame)

        return detections, frame